"""
任务协调Agent
"""

from typing import Dict, Any, List, Optional
from loguru import logger
from langchain_core.messages import BaseMessage, AIMessage, SystemMessage
from langchain_core.tools import BaseTool
from langchain_openai import ChatOpenAI

from ..core.config import settings


class OrchestratorAgent:
    """
    任务协调Agent，负责分解任务、调用工具并最终响应
    """
    
    def __init__(self, tools: List[BaseTool]):
        """
        初始化任务协调Agent
        
        Args:
            tools: 可供Agent使用的工具列表
        """
        # 初始化LLM
        llm = ChatOpenAI(
            model=settings.OPENAI_MODEL,
            api_key=settings.OPENAI_API_KEY,
            base_url=settings.OPENAI_BASE_URL,
            temperature=0.2, # 降低温度以获得更可预测、更稳定的工具使用
            streaming=True, # 依然可以流式传输，但在 process 中我们一次性获取结果
        )
        
        # 将工具绑定到LLM，这样LLM就知道有哪些工具可用
        self.llm = llm.bind_tools(tools)
        
        # 系统提示
        self.system_prompt = """你是一个专业的电商运营AI助手。
你的任务是分析用户的请求，并利用你所拥有的工具来完成这个请求。
逐步思考，规划你的行动。
首先，思考你是否需要使用工具。如果是，请调用你需要的工具。
在收到工具执行结果后，再次分析情况，决定下一步是调用另一个工具还是给出最终答复。
当你完成任务后，向用户提供一个简洁、清晰的最终答复。

当前平台信息: {platform_info}
"""
    
    async def process(
        self, messages: List[BaseMessage], platform_info: Optional[Dict[str, Any]] = None
    ) -> AIMessage:
        """
        处理用户消息 (非流式, 返回一个AIMessage)
        
        Args:
            messages: 消息历史 (LangChain格式)
            platform_info: 平台信息
            
        Returns:
            AIMessage: LLM的响应，可能包含工具调用或最终答案
        """
        try:
            # 格式化系统提示
            # 注意：我们不在每次调用时都创建新的 SystemMessage，
            # 而是将其作为invoke/ainvoke的第一个参数传入，这更高效。
            # 但为了简单起见，这里仍然构建消息列表。
            formatted_system_prompt = self.system_prompt.format(
                platform_info=platform_info or "未提供平台信息"
            )
            
            # 构建完整的消息历史
            chat_messages: List[BaseMessage] = [
                SystemMessage(content=formatted_system_prompt)
            ]
            chat_messages.extend(messages)
            
            # 调用LLM
            logger.info("Invoking OrchestratorAgent LLM with tools...")
            response = await self.llm.ainvoke(chat_messages)
            logger.info(f"LLM response: {response}")
            
            # 确保返回的是 AIMessage
            if not isinstance(response, AIMessage):
                raise TypeError(f"LLM did not return an AIMessage. Got {type(response)} instead.")

            return response
        
        except Exception as e:
            logger.error(f"OrchestratorAgent处理消息时出错: {str(e)}", exc_info=True)
            # 返回一个错误的 AIMessage，以便工作流可以捕获并处理
            return AIMessage(content=f"处理您的请求时出错: {str(e)}") 